import faiss
import pickle
import os
import numpy as np

# 加载文本嵌入
def load_embeddings(file_path):
    with open(file_path, 'rb') as f:
        return pickle.load(f)

# 构建 FAISS 索引
def build_faiss_index(embeddings):
    # 将嵌入转换为 numpy 数组
    embeddings = np.array(embeddings).astype(np.float32)
    
    # 创建 FAISS 索引
    dimension = embeddings.shape[1]  # 嵌入的维度
    
    # 使用倒排文件（IVF）索引来替代 IndexFlatL2
    nlist = 3  # 将数据分成多少个子集
    quantizer = faiss.IndexFlatL2(dimension)  # 基础的 L2 距离计算器
    index = faiss.IndexIVFFlat(quantizer, dimension, nlist, faiss.METRIC_L2)
    
    # 训练索引（只需要第一次训练）
    if not index.is_trained:
        print("训练索引...")
        index.train(embeddings)
    
    # 添加嵌入到索引中
    index.add(embeddings)
    
    return index

# 保存 FAISS 索引
def save_faiss_index(index, output_path):
    try:
        faiss.write_index(index, output_path)
        print(f"FAISS 索引已保存到 {output_path}")
    except Exception as e:
        print(f"保存索引时出现错误: {e}")

# 确保索引有效
def validate_faiss_index(index):
    if index.is_trained:
        print("索引已训练并且可以使用。")
    else:
        print("索引未训练，存在问题！")

def main():
    embeddings_path = "embeddings/knowledge_embeddings.pkl"
    faiss_index_path = "faiss/knowledge_index.index"
    
    # 确保 faiss 文件夹存在
    os.makedirs("faiss", exist_ok=True)
    
    # 加载嵌入
    embeddings = load_embeddings(embeddings_path)
    
    # 构建 FAISS 索引
    index = build_faiss_index(embeddings)
    
    # 验证索引
    validate_faiss_index(index)
    
    # 保存 FAISS 索引
    save_faiss_index(index, faiss_index_path)

if __name__ == "__main__":
    main()
